package com.shujia.custom.udaf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.io.DoubleWritable;
import java.util.List;

/**
 * UDAF 功能：将 ARRAY<STRING> 中的元素转换为小数后求和
 */
public class ArrayStringSum extends AbstractGenericUDAFResolver {

    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters) throws SemanticException {
        // 1. 校验输入参数类型
        if (parameters.length != 1) {
            throw new UDFArgumentTypeException(0, "必须提供1个参数");
        }
        String typeName = parameters[0].getTypeName();
        if (!typeName.startsWith("array<string>")) {
            throw new UDFArgumentTypeException(0, "输入必须是 ARRAY<STRING> 类型");
        }
        // 2. 返回计算器
        return new ArrayStringSumEvaluator();
    }

    /**
     * 核心计算逻辑实现
     */
    public static class ArrayStringSumEvaluator extends GenericUDAFEvaluator {
        // 类型检查器
        private ListObjectInspector arrayOI;
        private PrimitiveObjectInspector elementOI;
        private DoubleWritable result;

        @Override
        public ObjectInspector init(Mode mode, ObjectInspector[] parameters) throws HiveException {
            super.init(mode, parameters);

            // 初始化输入类型检查器
            if (mode == Mode.PARTIAL1 || mode == Mode.COMPLETE) {
                arrayOI = (ListObjectInspector) parameters[0];
                elementOI = (PrimitiveObjectInspector) arrayOI.getListElementObjectInspector();
            }

            // 输出类型为 Double
            result = new DoubleWritable(0.0);
            return PrimitiveObjectInspectorFactory.writableDoubleObjectInspector;
        }

        /**
         * 聚合缓冲区（存储中间结果）
         */
        static class SumAgg implements AggregationBuffer {
            double sum;
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer() {
            return new SumAgg();
        }

        @Override
        public void reset(AggregationBuffer agg) {
            ((SumAgg) agg).sum = 0.0;
        }

        /**
         * 核心逻辑：遍历数组元素并累加
         */
        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters) throws HiveException {
            if (parameters == null || parameters[0] == null) return;

            SumAgg sumAgg = (SumAgg) agg;
            List<?> array = arrayOI.getList(parameters[0]);

            for (Object element : array) {
                if (element != null) {
                    // 将字符串转换为 Double
                    String str = elementOI.getPrimitiveJavaObject(element).toString();
                    try {
                        double value = Double.parseDouble(str);
                        sumAgg.sum += value;
                    } catch (NumberFormatException e) {
                        // 非法字符串跳过（可根据需求调整）
                    }
                }
            }
        }

        /**
         * Map阶段返回部分结果
         */
        @Override
        public Object terminatePartial(AggregationBuffer agg) {
            return terminate(agg);
        }

        /**
         * Reduce阶段合并结果
         */
        @Override
        public void merge(AggregationBuffer agg, Object partial) {
            if (partial != null) {
                SumAgg sumAgg = (SumAgg) agg;
                sumAgg.sum += ((DoubleWritable) partial).get();
            }
        }

        /**
         * 返回最终结果
         */
        @Override
        public Object terminate(AggregationBuffer agg) {
            DoubleWritable output = new DoubleWritable(((SumAgg) agg).sum);
            return output;
        }
    }
}